_base_ = 'ssd300_coco.py'
# 默认24epoch
import os

# dataset settings
input_size = 300
# 修改类别数为2
model = dict(
    bbox_head=dict(
        type='SSDHead',
        num_classes=2,
    )
)

# 修改数据集相关配置（绝对路径）
dataset_type = 'CocoDataset'
data_root = '/root/mmdetection/dataset/SDB_9K_COCO/'
metainfo = {
    'classes': ('drone', 'bird'),
    'palette': [
        (253, 58, 52), (253, 159, 148),
    ]
}
# 三个loader
train_dataloader = dict(
    batch_size=4,
    num_workers=2,
    dataset=dict(
        dataset=dict(
            data_root=data_root,
            ann_file='train/annotations/SDB_9K_train.json',
            data_prefix=dict(img='train/images'),
        )
    )
)
val_dataloader = dict(
    batch_size=8,
    dataset=dict(
        pipeline={{_base_.test_pipeline}},
        dataset=dict(
            data_root=data_root,
            ann_file='val/annotations/SDB_9K_val.json',
            data_prefix=dict(img='val/images'),
            test_mode=True))
)
test_dataloader = dict(
    batch_size=8,
    dataset=dict(
        pipeline={{_base_.test_pipeline}},
        dataset=dict(
            data_root=data_root,
            ann_file='test/annotations/SDB_9K_test.json',
            data_prefix=dict(img='test/images'),
            test_mode=True))
)

# 修改评价指标相关配置
val_evaluator = dict(
    type='CocoMetric',
    ann_file=os.path.join(data_root, 'val/annotations/SDB_9K_val.json'),
    metric='bbox',
    format_only=False)
test_evaluator = dict(
    type='CocoMetric',
    ann_file=os.path.join(data_root, 'test/annotations/SDB_9K_test.json'),
    metric='bbox',
    format_only=False)

# optimizer
optim_wrapper = dict(
    type='OptimWrapper',
    optimizer=dict(type='SGD', lr=2e-3, momentum=0.9, weight_decay=5e-4))
load_from = 'https://download.openmmlab.com/mmdetection/v2.0/ssd/ssd300_coco/ssd300_coco_20210803_015428-d231a06e.pth'
